#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IR 控制流语句节点定义
"""

import ast
from dataclasses import dataclass, field
from typing import Optional, Iterator
from .base import IRStmt, IRStmtBlock
from ..base import IRNode, get_default_ir_node
from ..expressions import build_expr_from_ast, IRConstant
from ..expressions.basic import IRName


@dataclass
class IRReturn(IRStmt):
    """返回语句"""

    expr: Optional[IRNode] = field(default=None)
    node_type: str = "return"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        if self.expr is not None:
            yield self.expr
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.Return) -> None:
        """从 AST 返回节点构建 IR 返回节点内容

        Args:
            node: AST 返回节点
        """
        super().build_from_ast(node)

        if node.value:
            self.expr = build_expr_from_ast(node.value, self)
        else:
            self.expr = None


@dataclass
class IRBreak(IRStmt):
    """break 语句"""

    node_type: str = "break"

    def build_from_ast(self, node: ast.Break) -> None:
        """从 AST break 节点构建 IR break 节点内容

        Args:
            node: AST break 节点
        """
        super().build_from_ast(node)


@dataclass
class IRContinue(IRStmt):
    """continue 语句"""

    node_type: str = "continue"

    def build_from_ast(self, node: ast.Continue) -> None:
        """从 AST continue 节点构建 IR continue 节点内容

        Args:
            node: AST continue 节点
        """
        super().build_from_ast(node)


@dataclass
class IRIf(IRStmt):
    """条件语句"""

    test: IRNode = field(default_factory=get_default_ir_node)
    then_body: Optional[IRStmtBlock] = field(default=None)
    else_body: Optional[IRStmtBlock] = field(default=None)
    node_type: str = "if"

    def __post_init__(self):
        """初始化后处理"""
        if self.then_body is None:
            self.then_body = IRStmtBlock()

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.test
        if self.then_body:
            yield self.then_body
        if self.else_body:
            yield self.else_body
        yield from super().get_child_nodes()

    def get_local_symbols(self):
        """获取if语句中定义的所有本地符号

        Returns:
            本地符号的迭代器
        """
        # if语句本身不定义符号，但需要收集分支中的符号
        if self.then_body:
            yield from self.then_body.get_local_symbols()
        if self.else_body:
            yield from self.else_body.get_local_symbols()

    def build_from_ast(self, node: ast.If) -> None:
        """从 AST 条件节点构建 IR 条件节点内容

        Args:
            node: AST 条件节点
        """
        super().build_from_ast(node)

        self.test = build_expr_from_ast(node.test, self)

        # 构建 then 分支
        self.then_body = IRStmtBlock(parent=self)
        self.then_body.build_from_ast(node.body)

        # 构建 else 分支
        if node.orelse:
            self.else_body = IRStmtBlock(parent=self)
            self.else_body.build_from_ast(node.orelse)
        else:
            self.else_body = None


@dataclass
class IRForRange(IRStmt):
    """范围循环语句"""

    var: IRName = field(default_factory=lambda: IRName())  # 循环变量节点
    start: IRNode = field(default_factory=get_default_ir_node)
    end: IRNode = field(default_factory=get_default_ir_node)
    step: IRNode = field(default_factory=get_default_ir_node)
    body: Optional[IRStmtBlock] = field(default=None)
    node_type: str = "for_range"

    def __post_init__(self):
        """初始化后处理"""
        if self.body is None:
            self.body = IRStmtBlock()

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.var
        yield self.start
        yield self.end
        yield self.step
        if self.body:
            yield self.body
        yield from super().get_child_nodes()

    def get_local_symbols(self):
        """获取for循环中定义的所有本地符号

        Returns:
            本地符号的迭代器
        """

        # for循环定义循环变量符号
        # 使用变量节点的 get_symbol 方法创建符号
        yield self.var.get_symbol()

        # 收集循环体中的符号
        if self.body:
            yield from self.body.get_local_symbols()

    def build_from_ast(self, stmt: ast.For) -> bool:
        """从AST的For语句构建IRForRange节点

        Args:
            stmt: AST的For语句

        Returns:
            bool: 如果成功构建返回True，否则返回False
        """
        # 检查是否是range()调用
        super().build_from_ast(stmt)

        if not (
            isinstance(stmt.iter, ast.Call)
            and isinstance(stmt.iter.func, ast.Name)
            and stmt.iter.func.id == "range"
        ):
            return False

        # 设置循环变量
        # range() 循环只支持简单变量名，不支持元组解包
        if not isinstance(stmt.target, ast.Name):
            # 如果不是简单变量名，说明这不是合法的 range 循环
            # 例如：for (x, y) in range(10) 是不合法的
            return False

        # 创建循环变量的 IRName 节点
        var_node = IRName(id=stmt.target.id, parent=self)
        var_node.build_from_ast(stmt.target)
        self.var = var_node

        self.body = IRStmtBlock(parent=self)
        self.body.build_from_ast(stmt.body)

        # 处理range参数
        args = stmt.iter.args
        if len(args) == 1:
            # range(n) -> for (auto i = 0; i < n; ++i)
            end_ir = build_expr_from_ast(args[0], self)
            if end_ir is not None:
                self.start = IRConstant(value=0, parent=self)
                self.end = end_ir
                self.step = IRConstant(value=1, parent=self)
                self.line = getattr(stmt, "lineno", 0)
                self.source_line = self.line
                return True
        elif len(args) == 2:
            # range(start, end) -> for (auto i = start; i < end; ++i)
            start_ir = build_expr_from_ast(args[0], self)
            end_ir = build_expr_from_ast(args[1], self)
            if start_ir is not None and end_ir is not None:
                self.start = start_ir
                self.end = end_ir
                self.step = IRConstant(value=1, parent=self)
                self.line = getattr(stmt, "lineno", 0)
                self.source_line = self.line
                return True
        elif len(args) == 3:
            # range(start, end, step) -> for (auto i = start; i < end; i += step)
            start_ir = build_expr_from_ast(args[0], self)
            end_ir = build_expr_from_ast(args[1], self)
            step_ir = build_expr_from_ast(args[2], self)
            if start_ir is not None and end_ir is not None and step_ir is not None:
                self.start = start_ir
                self.end = end_ir
                self.step = step_ir
                self.line = getattr(stmt, "lineno", 0)
                self.source_line = self.line
                return True

        return False


@dataclass
class IRForEach(IRStmt):
    """遍历循环语句"""

    var: IRNode = field(default_factory=get_default_ir_node)
    iter_expr: IRNode = field(default_factory=get_default_ir_node)
    body: Optional[IRStmtBlock] = field(default=None)
    node_type: str = "for_each"

    def __post_init__(self):
        """初始化后处理"""
        if self.body is None:
            self.body = IRStmtBlock()

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.var
        yield self.iter_expr
        if self.body:
            yield self.body
        yield from super().get_child_nodes()

    def get_local_symbols(self):
        """获取for-each循环中定义的所有本地符号

        Returns:
            本地符号的迭代器
        """
        # for-each循环定义循环变量符号
        # 支持单变量和元组解包
        if isinstance(self.var, IRName):
            # 单变量：for x in ...
            yield self.var.get_symbol()
        else:
            # 元组或列表解包：for x, y in ... 或 for (x, y) in ...
            # 递归提取所有 IRName 节点
            for name_node in self._extract_names_from_target(self.var):
                yield name_node.get_symbol()

        # 收集循环体中的符号
        if self.body:
            yield from self.body.get_local_symbols()

    def _extract_names_from_target(self, target: IRNode) -> Iterator[IRName]:
        """从目标节点中提取所有 IRName 节点

        Args:
            target: 目标节点（可能是 IRName、IRTuple 或 IRList）

        Returns:
            IRName 节点的迭代器
        """
        if isinstance(target, IRName):
            yield target
        elif hasattr(target, "elements"):
            # IRTuple 或 IRList
            for element in getattr(target, "elements"):
                yield from self._extract_names_from_target(element)

    def build_from_ast(self, node: ast.For) -> None:
        """从 AST for 节点构建 IR for-each 节点内容

        Args:
            node: AST for 节点
        """
        super().build_from_ast(node)

        # 使用通用的表达式构建方法处理各种 target 类型
        # 支持：单变量 (x)、元组解包 (x, y)、列表解包 [x, y]
        self.var = build_expr_from_ast(node.target, self)

        self.iter_expr = build_expr_from_ast(node.iter, self)
        self.body = IRStmtBlock(parent=self)
        self.body.build_from_ast(node.body)


@dataclass
class IRWhile(IRStmt):
    """循环语句"""

    test: IRNode = field(default_factory=get_default_ir_node)
    body: Optional[IRStmtBlock] = field(default=None)
    node_type: str = "while"

    def __post_init__(self):
        """初始化后处理"""
        if self.body is None:
            self.body = IRStmtBlock()

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.test
        if self.body:
            yield self.body
        yield from super().get_child_nodes()

    def get_local_symbols(self):
        """获取while循环中定义的所有本地符号

        Returns:
            本地符号的迭代器
        """
        # while循环本身不定义符号，但需要收集循环体中的符号
        if self.body:
            yield from self.body.get_local_symbols()

    def build_from_ast(self, node: ast.While) -> None:
        """从 AST while 节点构建 IR while 节点内容

        Args:
            node: AST while 节点
        """
        super().build_from_ast(node)

        self.test = build_expr_from_ast(node.test, self)

        self.body = IRStmtBlock(parent=self)
        self.body.build_from_ast(node.body)
